[DE-8270] Model weights upload & download (SDK side) - #469
Open
luke-e-schaefer wants to merge 2 commits into
Open
[DE-8270] Model weights upload & download (SDK side)#469luke-e-schaefer wants to merge 2 commits into
luke-e-schaefer wants to merge 2 commits into
Conversation
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+149
to
+175
| transferred = 0 | ||
| finalized: List[Dict[str, Any]] = [] | ||
|
|
||
| def upload_part(part: dict) -> Dict[str, Any]: | ||
| nonlocal transferred | ||
| part_number = int(part[PART_NUMBER_KEY]) | ||
| offset = (part_number - 1) * part_size_bytes | ||
| chunk = _read_part(path, offset, part_size_bytes) | ||
| # Part PUTs are signed without the Content-Type condition, so they must | ||
| # be sent with no extra headers — including none of `requiredHeaders`. | ||
| etag = _put_bytes(part[URL_KEY], chunk) | ||
| if not etag: | ||
| raise RuntimeError( | ||
| f"Storage did not return an ETag for part {part_number}; " | ||
| "cannot finalize the multipart upload" | ||
| ) | ||
| transferred += len(chunk) | ||
| if on_progress is not None: | ||
| on_progress(min(transferred, total_bytes), total_bytes) | ||
| return {PART_NUMBER_KEY: part_number, ETAG_KEY: etag} | ||
|
|
||
| with ThreadPoolExecutor( | ||
| max_workers=min(CONCURRENT_PART_UPLOADS, len(parts)) | ||
| ) as pool: | ||
| finalized = list(pool.map(upload_part, parts)) | ||
|
|
||
| return sorted(finalized, key=lambda p: p[PART_NUMBER_KEY]) |
Contributor
There was a problem hiding this comment.
transferred counter
transferred += len(chunk) is a read-modify-write over three bytecodes (LOAD_DEREF → INPLACE_ADD → STORE_DEREF). CPython's GIL can switch threads between them — especially after the I/O-bound _put_bytes call releases it — so concurrent upload_part workers will race and silently drop updates. The on_progress callback will then report lower-than-actual byte counts, or even non-monotonically increasing values.
A threading.Lock keeps the update and the callback snapshot atomic:
import threading
def _upload_multipart(...):
transferred = 0
lock = threading.Lock()
def upload_part(part):
nonlocal transferred
...
etag = _put_bytes(part[URL_KEY], chunk)
if not etag:
raise RuntimeError(...)
with lock:
transferred += len(chunk)
current = min(transferred, total_bytes)
if on_progress is not None:
on_progress(current, total_bytes)
return {PART_NUMBER_KEY: part_number, ETAG_KEY: etag}Prompt To Fix With AI
This is a comment left during a code review.
Path: nucleus/model_weights.py
Line: 149-175
Comment:
**Race condition on `transferred` counter**
`transferred += len(chunk)` is a read-modify-write over three bytecodes (`LOAD_DEREF` → `INPLACE_ADD` → `STORE_DEREF`). CPython's GIL can switch threads between them — especially after the I/O-bound `_put_bytes` call releases it — so concurrent `upload_part` workers will race and silently drop updates. The `on_progress` callback will then report lower-than-actual byte counts, or even non-monotonically increasing values.
A `threading.Lock` keeps the update and the callback snapshot atomic:
```python
import threading
def _upload_multipart(...):
transferred = 0
lock = threading.Lock()
def upload_part(part):
nonlocal transferred
...
etag = _put_bytes(part[URL_KEY], chunk)
if not etag:
raise RuntimeError(...)
with lock:
transferred += len(chunk)
current = min(transferred, total_bytes)
if on_progress is not None:
on_progress(current, total_bytes)
return {PART_NUMBER_KEY: part_number, ETAG_KEY: etag}
```
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
Comment on lines
+128
to
+138
| def _upload_single( | ||
| path: str, | ||
| upload_url: str, | ||
| headers: Dict[str, str], | ||
| total_bytes: int, | ||
| on_progress: Optional[ProgressCallback], | ||
| ) -> None: | ||
| with open(path, "rb") as handle: | ||
| _put_bytes(upload_url, handle, headers) | ||
| if on_progress is not None: | ||
| on_progress(total_bytes, total_bytes) |
Contributor
There was a problem hiding this comment.
For a single PUT (files up to 5 GB), _upload_single passes an open file handle directly to requests.put and then calls on_progress(total_bytes, total_bytes) once the whole transfer finishes. Callers expecting incremental updates — e.g., a progress bar — will see nothing for minutes, then a sudden jump to 100 %. Multipart uploads (≥ 5 GB) do get per-part callbacks, so the behaviour differs across the size boundary. Consider passing a generator or read-wrapped file object to _put_bytes to produce intermediate callbacks, or at least document the single-shot behaviour on the callback parameter.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nucleus/model_weights.py
Line: 128-138
Comment:
**Single-part progress callback fires only on completion**
For a single PUT (files up to 5 GB), `_upload_single` passes an open file handle directly to `requests.put` and then calls `on_progress(total_bytes, total_bytes)` once the whole transfer finishes. Callers expecting incremental updates — e.g., a progress bar — will see nothing for minutes, then a sudden jump to 100 %. Multipart uploads (≥ 5 GB) do get per-part callbacks, so the behaviour differs across the size boundary. Consider passing a generator or read-wrapped file object to `_put_bytes` to produce intermediate callbacks, or at least document the single-shot behaviour on the callback parameter.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Base automatically changed from
update-nuc-sdk-for-new-eval-stuff-pt1
to
master
August 11, 2026 14:24
Mirrors the REST surface shipped in scaleapi#149063 so users can attach a weights artifact to a model and fetch it back from Python. - `NucleusClient.upload_model_weights` / `download_model_weights` — the two methods the server PR's API-docs pages already document — plus `get_model_weights` and `delete_model_weights` for the remaining routes. - `Model.upload_weights()` / `download_weights()` / `weights()` / `delete_weights()` delegate to the client, matching how `Benchmark` does it. - New `ModelWeights` metadata type parsed from the weights DTO. The weights routes serialize camelCase both ways, unlike most of this SDK's endpoints, so the new payload keys are grouped and labelled in `constants.py`. - Transfers go straight to storage via presigned URLs and never through the API, so artifacts aren't subject to API request-size limits. Over 5 GB the server hands back multipart parts, which upload 4 at a time; `on_progress` reports `(bytes_transferred, total_bytes)`. - Size is checked against the server's 10 GB cap before presign, so an oversized file fails without a network round-trip. Two things worth knowing for review: part PUTs must be sent with *no* headers (they're signed without the Content-Type condition, so forwarding `requiredHeaders` makes S3 reject the signature), and download resolves the signed URL via `?json=1` rather than following the 302, so the API's auth headers are never sent to storage. 24 mock-based unit tests in `tests/test_model_weights.py`; version bumped to 0.19.1 (additive, per CLAUDE.md). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The docstrings are what users read, so they shouldn't describe how the artifact gets stored or moved. Dropped the presign/multipart/direct-to-storage narration from every public docstring, the `ModelWeights` attribute docs, and the CHANGELOG, leaving what a caller actually needs: what the method does, who can call it, the size limit, and the arguments. Also made the transfer helpers private (`_presign_payload`, `_transfer_weights_to_storage`, `_stream_weights_to_file`, `_finalize_payload`) so the mechanics don't show up in the generated API docs at all, rather than only being reworded. Kept the two in-body comments that explain why part uploads send no headers and why the download URL is fetched as JSON — those aren't user-visible and each one guards a real footgun. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
luke-e-schaefer
force-pushed
the
lukeschaefer/de-8270-upload-download-model-weights
branch
from
August 11, 2026 14:35
df6ede7 to
83c25f7
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
SDK side of model weights upload / download, mirroring the REST surface merged in scaleapi#149063.
The two primary methods are the ones the server PR's API-docs pages (
ApiDocsPage/models-python/{upload,download}-model-weights.md) already document, so the published docs and the SDK agree:Added
NucleusClient.upload_model_weights(model, path, *, content_type=None, original_filename=None, checksum_sha256=None, on_progress=None)— presign → PUT direct to storage → finalize. ReturnsModelWeights.NucleusClient.download_model_weights(model, path, *, on_progress=None)— resolves the signed URL and streams to disk (creating parent dirs). Returns the path written.get_model_weights(model)/delete_model_weights(model)for the remaining two routes.Model.upload_weights()/.download_weights()/.weights()/.delete_weights()— thin delegation to the client, matching howBenchmarkwraps its client methods.ModelWeightsmetadata type:present,status,size_bytes,original_filename,content_type,download_url.All model arguments accept either a
Modelor a bare model id (prj_*).Notes for review
requiredHeaders(which the single PUT does need) makes S3 reject the signature. There's a test pinning this in both directions — it's the easiest thing to get wrong here, and the frontend hook in 149063 has the same split.?json=1to fetch the signed URL rather than following the 302, so the API's auth headers are never sent to storage.constants.py.Tests / Version
tests/test_model_weights.py— 24 mock-based unit tests (no live API, no real S3): DTO parsing, payload builders, single vs. multipart transfer, header split, ETag/failure handling, progress callbacks, download streaming, all four client methods, and theModelwrappers.pylint nucleus10.00/10,mypy --ignore-missing-imports nucleusclean, ruff clean, black + isort clean, 61 mock-based tests passing across the eval/benchmark/leaderboard/weights suites.pyproject.toml→ 0.19.1 + CHANGELOG entry (patch bump: additive new methods, perCLAUDE.md).resolves https://linear.app/scale-epd/issue/DE-8270
🤖 Generated with Claude Code
Greptile Summary
Adds model weights upload and download to the Nucleus Python SDK — a presign → direct-to-S3 PUT → finalize upload flow (with multipart support above 5 GB and 4 concurrent part workers), a streaming download via a signed URL fetched as JSON, and the supporting
ModelWeightsDTO and fourNucleusClient/Modelconvenience methods.model_weights.pyintroduces all transfer machinery: single-PUT and multipart upload paths (correctly omittingrequiredHeadersfrom part PUTs to avoid S3 signature rejection), chunked streaming download withos.makedirsfor missing parents, andProgressCallbacksupport.__init__.pywraps these intoupload_model_weights,download_model_weights,get_model_weights, anddelete_model_weights, with a client-side 10 GB guard before the presign round-trip.model.pyadds thin delegation wrappers (upload_weights,download_weights,weights,delete_weights) following the same pattern used byBenchmark.Confidence Score: 5/5
Important Files Changed
Sequence Diagram
sequenceDiagram participant Caller participant NucleusClient participant NucleusAPI participant S3 Note over Caller,S3: Upload flow Caller->>NucleusClient: upload_model_weights(model, path) NucleusClient->>NucleusClient: size check (≤10 GB) NucleusClient->>NucleusAPI: "POST model/{id}/weights/presign" NucleusAPI-->>NucleusClient: "{uploadId, uploadUrl|parts, requiredHeaders}" alt "Single PUT (< 5 GB)" NucleusClient->>S3: PUT file bytes (with requiredHeaders) S3-->>NucleusClient: ETag else Multipart (≥ 5 GB) par 4 concurrent workers NucleusClient->>S3: PUT part 1 (no headers) NucleusClient->>S3: PUT part 2 (no headers) NucleusClient->>S3: PUT part N (no headers) end S3-->>NucleusClient: ETags per part end NucleusClient->>NucleusAPI: "POST model/{id}/weights/finalize" NucleusAPI-->>NucleusClient: ModelWeights metadata NucleusClient-->>Caller: ModelWeights Note over Caller,S3: Download flow Caller->>NucleusClient: download_model_weights(model, path) NucleusClient->>NucleusAPI: "GET model/{id}/weights/download?json=1" NucleusAPI-->>NucleusClient: "{url: signed-get-url}" NucleusClient->>S3: GET signed URL (streaming) S3-->>NucleusClient: file bytes (chunked) NucleusClient-->>Caller: path writtenReviews (4): Last reviewed commit: "Make the weights docstrings user-facing" | Re-trigger Greptile